# Ref: Cross section database for carbon atoms and ions: Electron impact ionization
#      excitation, and charge exchange in collisions with hydrogen atoms
# Journal: Atomic data and nuclear data tables 92 (2006) 407-455
import numpy as np
import math
import matplotlib.pyplot as plt


class CrossSection:
    def __init__(self, threshold_energy, reaction_name, is_ground_state = True):
        self.threshold_energy = threshold_energy

        data_temp = [[self.threshold_energy],[0.0]]
        for energy in np.arange(float(int(self.threshold_energy+1)), 100.0, 1.0):
            CS = 0.0
            data_temp[0].append(energy)
            data_temp[1].append(CS)

        for energy in np.arange(100.0, 1000.0, 10.0):
            CS = 0.0
            data_temp[0].append(energy)
            data_temp[1].append(CS)

        data_np = np.asarray(data_temp)
        data_np = np.transpose(data_np)

        #cs_data: cross section data
        self.cs_data = data_np.copy()
        self.collision_strength_data = data_np.copy()

        self.reaction_name = reaction_name

        self.is_ground_state = is_ground_state

    def get_cross_section(self, energy):
        cs_temp = 0.0
        if energy <= self.threshold_energy:
            return 0.0
        elif energy >= self.cs_data[-1,0]:
            return self.cs_data[-1,1]
        else:
            for i in range(2,self.cs_data.shape[0]-1):
                if self.cs_data[i+1,0] >= energy:
                    e0 = self.cs_data[i,0]
                    e1 = self.cs_data[i+1,0]
                    cs0 = self.cs_data[i,1]
                    cs1 = self.cs_data[i+1,1]
                    c0 = (e1-energy) / (e1-e0)
                    c1 = (energy-e0) / (e1-e0)
                    return cs0 * c0 + cs1 * c1
            print("Cross Section Error: get_cross_section")

class CS_Excitation:
    def __init__(self, species_name):
        self.equation_number = 0
        self.a0 = 5.2917706e-11
        self.max_CS = 2.0e-22
        self.Vif = 0.0
        self.A = 0.0
        self.B = 0.0
        self.C = 0.0
        self.D = 0.0
        self.E = 0.0
        self.F = 0.0

        self.cs_list = []
        self.cs_total = None

        self.data_prefix = "data/"
        self.species_name = species_name

        self.min_Vif = 1000.0


    # Equation (6)
    def excitation6(self, Te):
        X = Te / self.Vif
        Qif =  self.A + self.B / X + self.C / (X*X) + self.D / (X*X*X) + self.E * math.log(X)
        #print("excitation6: ", self.A, self.B, self.C, self.D, self.E)
        return Qif

    # Equation (7)
    def excitation7(self, Te):
        #print(A, B, C)
        X = Te / self.Vif
        Qif =  self.A / (X*X) + self.B * math.exp( - self.F * X ) + self.C * math.exp(- 2.0 * self.F * X) + self.D * math.exp( - 3.0 * self.F * X ) + self.E * math.exp( - 4.0 * self.F * X )
        return Qif

    # Equation (15)
    def excitation15(self, Te):
        X = Te / self.Vif
        Qif =  self.P * X + self.Q
        return Qif

    # equation (3)
    def get_cross_section(self, collision_strength, Te, omega):
        return 1.1969e-19 * collision_strength / ( omega * Te )

    #omega: the statistical weight of the initial state, alse called multiplet/degenerate 
    #简并度，多重态
    def add_cross_section(self, id, Vif, A, B, C, D, E, F, omega, equation_number, reaction_name, X1=0.0, P=0.0, Q=0.0):
        self.equation_number = equation_number
        self.Vif = Vif 
        self.A = A
        self.B = B
        self.C = C
        self.D = D
        self.E = E
        self.F = F
        self.X1 = X1
        self.P = P
        self.Q = Q
        
        self.omega = omega
        
        if self.Vif < self.min_Vif:
            self.min_Vif = self.Vif

        cs = CrossSection(self.Vif, reaction_name)
        for i in range(1, cs.cs_data.shape[0]):
            energy = cs.cs_data[i,0]

            if self.X1 > 0.0 and energy / self.Vif < self.X1:
                collision_strength_i = self.excitation15(energy)
            else:
                if self.equation_number == 6:
                    collision_strength_i = self.excitation6(energy)
                elif self.equation_number == 7:
                    collision_strength_i = self.excitation7(energy)

            if collision_strength_i < 0.0:
                collision_strength_i = 0.0

            cs.collision_strength_data[i,0] = energy / self.Vif
            cs.collision_strength_data[i,1] = collision_strength_i


            cs_i = self.get_cross_section(collision_strength_i, energy, self.omega)
            if cs_i < 0.0:
                cs_i = 0.0

            cs.cs_data[i,0] = energy
            cs.cs_data[i,1] = cs_i

        self.cs_list.append(cs)


    
              
    def plot_cs(self):
        for i in range(len(self.cs_list)):
            #if "Excitation_C_2s22p23P-2s2p33P" == self.cs_list[i].reaction_name:
            #    plt.semilogx(self.cs_list[i].cs_data[:,0], self.cs_list[i].cs_data[:,1])
            plt.semilogx(self.cs_list[i].cs_data[:,0], self.cs_list[i].cs_data[:,1])
            if self.cs_list[i].cs_data[:,1].min() < 0.0:
                print("cs below zero: ", self.cs_list[i].reaction_name, self.cs_list[i].cs_data[:,1].min())
            #for j in range(self.cs_list[i].cs_data[:,0].shape[0]):
                #print("cs energy: ", self.cs_list[i].cs_data[j,0], self.cs_list[i].collision_strength_data[j,0])
            #    print("cs energy: ", self.cs_list[i].cs_data[j,0], self.cs_list[i].cs_data[j,1])
        plt.xlim((0.0, 500))
        #plt.ylim((0.0, 10))
        #plt.legend()
        plt.savefig("fig/Excitation_" + self.species_name + ".svg")

    def save_cs(self):
        for i in range(len(self.cs_list)):
            np.savetxt(self.data_prefix + self.cs_list[i].reaction_name + "_ref.dat", self.cs_list[i].cs_data, fmt='%1.5e')


    def cal_total_cs(self):
        self.cs_total = CrossSection(self.min_Vif, "Excitation_" + self.species_name + "_total")
        for i in range(1, self.cs_total.cs_data.shape[0]):
            energy = self.cs_total.cs_data[i,0]
            cs_i = 0.0
            for j in range(len(self.cs_list)):
                cs_i = cs_i + self.cs_list[j].threshold_energy * self.cs_list[j].get_cross_section(energy)
            cs_i = cs_i / self.min_Vif
            self.cs_total.cs_data[i,1] = cs_i
        np.savetxt(self.data_prefix + self.cs_total.reaction_name + ".dat", self.cs_total.cs_data, fmt='%1.5e')

        plt.clf()
        plt.semilogx(self.cs_total.cs_data[:,0], self.cs_total.cs_data[:,1])
        plt.xlim((0.0, 1000))
        plt.savefig("fig/Excitation_" + self.species_name + "_total.svg")

if __name__ == '__main__':
    cs = CS_Excitation()
    cs.generate_cross_section()
    cs.plot_cs()
    cs.save_cs()
    cs.cal_total_cs()

